import paddle.fluid as fluid
import paddle.fluid.layers as L

__all__ = ['CrossEntropyLoss']


def cross_entropy(input, target, label_smooth=None):
    num_classes = input.shape[1]
    target = L.squeeze(target, axes=[1])
    one_hot_target = fluid.one_hot(input=target, depth=num_classes)
    if label_smooth:
        target = L.label_smooth(one_hot_target, epsilon=label_smooth)
    loss = L.softmax_with_cross_entropy(input, target, soft_label=True)
    loss = L.mean(loss)
    return loss


class CrossEntropyLoss:

    def __init__(self, label_smooth=None):
        self.label_smooth = label_smooth

    def __call__(self, input, target):
        if target.shape[1] == 1 and not self.label_smooth:
            loss = L.softmax_with_cross_entropy(input, target, soft_label=False)
            return L.mean(loss)

        if target.shape[1] == 1:
            num_classes = input.shape[1]
            target = L.squeeze(target, axes=[1])
            target = fluid.one_hot(input=target, depth=num_classes)
        if self.label_smooth:
            target = L.label_smooth(target, epsilon=self.label_smooth)
        loss = L.softmax_with_cross_entropy(input, target, soft_label=True)
        return L.mean(loss)